from fastmcp import FastMCP
import numpy as np

# 创建一个MCP服务器实例
mcp = FastMCP("Matrix Operations")

def matrix_add(a: list, b: list) -> list:
    """矩阵加法"""
    return (np.array(a) + np.array(b)).tolist()

def matrix_subtract(a: list, b: list) -> list:
    """矩阵减法"""
    return (np.array(a) - np.array(b)).tolist()

def matrix_multiply(a: list, b: list) -> list:
    """矩阵乘法"""
    return (np.dot(np.array(a), np.array(b))).tolist()

def matrix_divide(a: list, b: list) -> list:
    """矩阵元素除法"""
    # 注意：这里假设b中的每个元素都不为0
    return (np.divide(np.array(a), np.array(b))).tolist()

# 添加工具到MCP Server
@mcp.tool()
def add_tool(a: list, b: list) -> list:
    return matrix_add(a, b)

@mcp.tool()
def subtract_tool(a: list, b: list) -> list:
    return matrix_subtract(a, b)

@mcp.tool()
def multiply_tool(a: list, b: list) -> list:
    return matrix_multiply(a, b)

@mcp.tool()
def divide_tool(a: list, b: list) -> list:
    return matrix_divide(a, b)

if __name__ == "__main__":
    mcp.run(transport='stdio')